const path = require('path');
const fs = require('fs');
const http = require('http');
const https = require('https');
const dgram = require('dgram');
const dns2 = require('dns2');
const sqlite3 = require('sqlite3');

const { Packet } = dns2;

/*
 * 所有的Resource Record (RR) TYPEs
 * https://www.iana.org/assignments/dns-parameters/dns-parameters.xhtml#dns-parameters-4
*/
Packet.TYPE['DS'] = 43; // Delegation Signer	[RFC4034]
Packet.TYPE['RRSIG'] = 46; // RRSIG	[RFC4034]
Packet.TYPE['NSEC'] = 47; // NSEC	[RFC4034][RFC9077]
Packet.TYPE['DNSKEY'] = 48; // DNSKEY	[RFC4034]

function getRandomInt(min, max) {
  min = Math.ceil(min); // 向上取整，确保是整数
  max = Math.floor(max); // 向下取整，确保是整数
  return Math.floor(Math.random() * (max - min + 1)) + min;
}

const my_request_id = getRandomInt(0, 65535);
console.log('my_request_id', my_request_id);

function getDateStr() {
  //获取本地时间字符串
  return (new Date()).toLocaleString();
}

const SqliteDb = (dbFile) => {
  const mkdirsSync = function (dirname) {
    if (fs.existsSync(dirname)) {
      return true;
    }
    if (mkdirsSync(path.dirname(dirname))) {
      fs.mkdirSync(dirname);
      return true;
    }
  };
  mkdirsSync(path.dirname(dbFile));
  const pri = {};
  pri.db = new sqlite3.Database(dbFile);

  const pub = {};

  pub.fStmtDo = (stmt, type = 'all', v = []) => {
    return new Promise((resolve, reject) => {
      v.push((err, rows) => {
        if (err) {
          reject(err);
          return;
        }
        resolve(rows);
      });
      stmt[type](...v);
    });
  };

  pub.fDoDb = (type, sql) => {
    // 改成异步
    return new Promise((resolve, reject) => {
      pri.db[type](sql, (err, rows) => {
        if (err) {
          reject(err);
          return;
        }
        resolve(rows);
      });
    });
  }

  pub.fCreateTable = async (tableName) => {
    // 建表
    const sql = `CREATE TABLE "${tableName}"(
      "id" INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT,
      "created" BIGINT NOT NULL,
      "updated" BIGINT NOT NULL,
      "name" TEXT NOT NULL COLLATE NOCASE,
      "type" INTEGER NOT NULL,
      "value" TEXT NOT NULL,
      CONSTRAINT "data_unique" UNIQUE ("name" COLLATE NOCASE ASC, "type" ASC)
    )`;
    try {
      return await pub.fDoDb('run', sql);
    } catch (ex) {
      return ex;
    }
  };

  pub.begin = () => {
    return pub.fDoDb('run', "BEGIN TRANSACTION");
  };

  pub.commit = () => {
    return pub.fDoDb('run', "COMMIT");
  };

  pub.fSelect = async (sTableName, mWhere = {}, iLimit = 10) => {
    // 查找数据
    const aWhere = [];
    const aValues = [];
    for (const k in mWhere) {
      const v = mWhere[k];
      aWhere.push(v.sql ? v.sql : `"${k}"=?`);
      aValues.push(v.value ? v.value : v);
    }
    const sWhere = aWhere ? ' WHERE (' + aWhere.join(')AND(') + ')' : '';
    const sql = `SELECT * FROM "${sTableName}"${sWhere} LIMIT ${iLimit}`;
    const stmt = pri.db.prepare(sql);
    return await pub.fStmtDo(stmt, 'all', aValues);
  };

  pub.fUpsert = async (sTableName, mWhere, mInsert, mUpdate) => {
    const aValues = [];
    const aInsertKeys = [];
    for (const sKey in mWhere) {
      aInsertKeys.push(sKey);
      aValues.push(mWhere[sKey]);
    }
    for (const sKey in mInsert) {
      aInsertKeys.push(sKey);
      aValues.push(mInsert[sKey]);
    }
    for (const sKey in mUpdate) {
      aValues.push(mUpdate[sKey]);
    }
    const aSql = [];
    const sInsertKeys = aInsertKeys.join(',');
    aSql.push(`INSERT INTO "${sTableName}" (${sInsertKeys})`);
    const sInsertValues = Array(aInsertKeys.length).fill('?');
    aSql.push(`VALUES (${sInsertValues})`);
    const sWhereKeys = Object.keys(mWhere).join(',');
    aSql.push(`ON CONFLICT (${sWhereKeys})`);
    const aUpdateSets = [];
    for (const k in mUpdate) {
      aUpdateSets.push(`${k}=?`);
    }
    const sUpdateSets = aUpdateSets.join(',');
    aSql.push(`DO UPDATE SET ${sUpdateSets};`);
    const sSql = aSql.join(' ');
    const stmt = pri.db.prepare(sSql);
    await pub.fStmtDo(stmt, 'run', aValues);
  };

  pub.fPutRecord = async (sTableName, mZoneData, useTransaction = false) => {
    if (useTransaction) {
      await pub.begin();
    }
    for (const name in mZoneData) {
      const mZoneType = mZoneData[name];
      for (const sType in mZoneType) {
        const aValue = mZoneType[sType];
        const sValue = JSON.stringify(aValue);
        const mWhere = {};
        mWhere.name = name;
        mWhere.type = sType;
        const mInsert = {};
        mInsert.created = +new Date();
        mInsert.updated = +new Date();
        mInsert.value = sValue;
        const mUpdate = {};
        mUpdate.updated = +new Date();
        mUpdate.value = sValue;
        await pub.fUpsert(sTableName, mWhere, mInsert, mUpdate);
      }
    }
    if (useTransaction) {
      await pub.commit();
    }
  };

  pub.fGetRecord = async (sTableName, name, type) => {
    const mWhere = {};
    mWhere.name = name;
    mWhere.type = type;
    if (sTableName === 'zone_cache') {
      mWhere.updated = {};
      mWhere.updated.sql = 'updated>?';
      mWhere.updated.value = (+new Date()) - 1000 * 60;
    }
    const aRows = await pub.fSelect(sTableName, mWhere, 1);
    if (aRows && aRows.length) {
      const mRow = aRows[0];
      return JSON.parse(mRow.value);
    }
  };

  return pub;
};

const mkdirsSync = (dirname) => {
  // 递归创建文件夹
  if (fs.existsSync(dirname)) {
    return true;
  }
  if (mkdirsSync(path.dirname(dirname))) {
    fs.mkdirSync(dirname);
    return true;
  }
};

const writeFileSync = (filePath, data) => {
  mkdirsSync(path.dirname(filePath));
  fs.writeFileSync(filePath, data);
};

function download(url, dest) {
  console.log(new Date(), `downloading ${url}`);
  const downloadProtocol = url.startsWith('https') ? https : http;
  return new Promise((resolve, reject) => {
    const oUrl = new URL(url);
    const options = {
      method: 'GET',
      hostname: oUrl.hostname,
      path: oUrl.pathname + oUrl.search,
    };
    if (oUrl.port) {
      options.port = oUrl.port;
    }
    options.headers = {};
    options.headers["user-agent"] = "curl/7.64.1";
    downloadProtocol.get(options, (response) => {
      mkdirsSync(path.dirname(dest));
      const fileStream = fs.createWriteStream(dest);
      response.pipe(fileStream);
      fileStream.on('finish', () => {
        fileStream.close();
        resolve(dest);
      });
    }).on('error', (err) => {
      fs.unlink(dest, () => {
        reject(err);
      });
    });
  });
}

const getRootZoneFileData = async (zoneFilePath) => {
  const url_root_zone = 'https://www.internic.net/domain/root.zone';
  if (fs.existsSync(zoneFilePath)) {
    return fs.readFileSync(zoneFilePath);
  }
  await download(url_root_zone, zoneFilePath);
  return fs.readFileSync(zoneFilePath);
};

const parseZoneData = (sZoneFileData) => {
  if (typeof (sZoneFileData) !== 'string') {
    sZoneFileData = sZoneFileData.toString();
  }
  const aZoneData = [];
  const aLines = sZoneFileData.split("\n");
  for (const sLine of aLines) {
    if (sLine === '') {
      continue;
    }
    const aRow = sLine.split(/\t+/);
    aZoneData.push(aRow);
  }
  return aZoneData;
};


/**
 * 从字符串两端删除非法字符
 *
 * @param {string} str - 待处理的字符串
 * @returns {string} - 删除非法字符后的字符串
 */
const trimIllegalCharacters = (str) => {
  // 匹配所有不是小写字母和数字的字符（从字符串开头或结尾），并用空字符串替换它们
  const illegalCharactersRegExp = /^[^a-z\d_]+|[^a-z\d_]+$/ig;
  const sRet = str.replace(illegalCharactersRegExp, '');
  return sRet.toLowerCase();
};

const fromZoneDataToMap = (aZoneData) => {
  // 将 parseZoneData() 的结果再转换成可快速查询的Map类型
  const mRet = {};
  for (const aRow of aZoneData) {
    const sName = trimIllegalCharacters(aRow[0]); // 域名
    const sType = aRow[3]; // 记录类型，如: NS, A, AAAA
    const iType = typeof (sType) === 'string' ? Packet.TYPE[sType] : sType;
    if (!iType) {
      console.log('type is not found =>', aRow);
      continue;
    }
    const sValue = aRow[4]; // 解析记录值
    if (!mRet[sName]) {
      mRet[sName] = {};
    }
    if (!mRet[sName][iType]) {
      mRet[sName][iType] = [];
    }
    mRet[sName][iType].push(sValue);
  }
  return mRet;
};

const pushEDNS = (request, clientIp) => {
  if (!clientIp) {
    //clientIp = '5.28.39.0/24';
    clientIp = "114.114.114.114";
  }
  if (clientIp) {
    request.additionals.push(Packet.Resource.EDNS([
      Packet.Resource.EDNS.ECS(clientIp),
    ]));
  }
};

const UDPClient = (msg, dnsAddr) => {
  if (!dnsAddr) {
    //dnsAddr = '1.1.1.1';
  }
  return new Promise((resolve, reject) => {
    const client = dgram.createSocket('udp4');
    client.on('error', (err) => {
      client.close();
      reject(`[${getDateStr()}] [ERROR] - ${err}`);
    });
    client.on('message', (fbMsg, fbRinfo) => {
      resolve(fbMsg);
      client.close();
    });
    client.send(msg, 53, dnsAddr, (err) => {
      if (err) {
        client.close();
        reject(`[${getDateStr()}] [ERROR] - ${err}`);
      }
    });
  });
};

const getPacket = (name, type = 1) => {
  const request = new Packet();
  request.header.id = my_request_id;
  request.header.rd = 1; // recursive
  request.questions.push({
    name,
    class: Packet.CLASS.IN,
    type,
  });
  return request;
};

const randomArrayElement = (arr) => {
  // 从数组中随机选择一个索引
  const randomIndex = Math.floor(Math.random() * arr.length);
  // 返回该索引对应的元素
  return arr[randomIndex];
}

function bin2hex(binaryData) {
  const buffer = Buffer.from(binaryData);
  return buffer.toString('hex');
}

function base64Encode(data) {
  const buffer = Buffer.from(data);
  return buffer.toString('base64');
}

function base64Decode(encodedData) {
  return Buffer.from(encodedData, 'base64');
}

const fGetValueOfZone = (item) => {
  if (item.type === Packet.TYPE.CNAME && item.domain) {
    return item.domain;
  }
  if (item.type === Packet.TYPE.SOA && item.primary) {
    return item.primary;
  }
  if (item.type === Packet.TYPE.NS && item.ns) {
    return item.ns;
  }
  if (item.type === Packet.TYPE.MX && item.exchange) {
    return [item.exchange, item.priority];
  }
  if (item.type === Packet.TYPE.TXT && item.data) {
    return item.data;
  }
  if (item.address) {
    if (item.address === '0.0.0.1') {
      //return;
    }
    return item.address;
  }
  if (item.data) {
    return base64Encode(item.data);
  }
  if (Object.keys(item).length > 3) {
    console.log('whats the value in item of zone? => ', item);
  }
};

const fromResponseToArrayZoneData = (response) => {
  const aData = [];
  for (const k in response) {
    if (k === 'header' || k === 'requetions') {
      continue;
    }
    if (k === 'toBuffer' || k === 'toBase64URL') {
      continue;
    }
    if (k === 'recursive') {
      continue;
    }
    const items = response[k];
    if (!Array.isArray(items)) {
      console.log('not find the k of response, k =>', k, ', items =>', items);
      continue;
    }
    for (const item of items) {
      if (!item.name) {
        continue;
      }
      if (item.type === 0) {
        continue;
      }
      const sValue = fGetValueOfZone(item);
      if (!sValue) {
        continue;
      }
      const aZone = [];
      aZone.push(item.name);
      aZone.push(item.ttl);
      aZone.push(item.class);
      aZone.push(item.type);
      aZone.push(sValue);
      aData.push(aZone);
    }
  }
  return aData;
};


const cf_custom_hostname = (mZoneData, hostname) => {
  mZoneData[hostname] = {};
  mZoneData[hostname][Packet.TYPE.A] = ['104.18.3.161', '104.18.2.161'];
  mZoneData[hostname][Packet.TYPE.AAAA] = ['2606:4700::6812:3a1', '2606:4700::6812:2a1'];
  mZoneData[hostname][Packet.TYPE.CNAME] = ['ss.feieryun.net'];
  mZoneData[`_acme-challenge.${hostname}`] = {};
  mZoneData[`_acme-challenge.${hostname}`][Packet.TYPE.TXT] = ['Ky1DBRt016i7EAcD13rqywPjmqM5t9IgCYz-llRqj7A'];
  mZoneData[`_acme-challenge.${hostname}`][Packet.TYPE.CNAME] = [`${hostname}.079c2023c61792ee.dcv.cloudflare.com`];
  mZoneData[`_cf-custom-hostname.${hostname}`] = {};
  mZoneData[`_cf-custom-hostname.${hostname}`][Packet.TYPE.TXT] = ['cc425838-6d25-47ee-b899-bfe76db2cc6d'];
};

const main = async () => {

  const iDnsAddrType = Packet.TYPE.A;

  const sRootZoneData = await getRootZoneFileData('domain-root.zone');
  const aRootZoneData = parseZoneData(sRootZoneData);
  const mRootZoneData = fromZoneDataToMap(aRootZoneData);

  const dbFile = 'domain.sqlite3.db';
  const db = SqliteDb(dbFile);
  const err = await db.fCreateTable('zone_root');
  if (!err) {
    console.log('正在将 [domain-root.zone] 导入到 [zone_root] 表 ...');
    await db.fPutRecord('zone_root', mRootZoneData, true);
    console.log('已成功导入 [zone_root] 表');
  }
  await db.fCreateTable('zone_cache');
  await db.fCreateTable('zone_data');

  const mZoneData = {};
  cf_custom_hostname(mZoneData, 'a.wangke.fyi');
  cf_custom_hostname(mZoneData, 'vip.waruandaigua.xyz');
  cf_custom_hostname(mZoneData, 'a.shanggpppda.shop');

  await db.fPutRecord('zone_data', mZoneData, true);

  const getAllRecordValueInDb = async (name, type) => {
    const aRecordCache = await db.fGetRecord('zone_cache', name, type);
    if (aRecordCache && aRecordCache.length) {
      return aRecordCache;
    }

    const aRecordData = await db.fGetRecord('zone_data', name, type);
    if (aRecordData && aRecordData.length) {
      return aRecordData;
    }

    const aRecordRoot = await db.fGetRecord('zone_root', name, type);
    if (aRecordRoot && aRecordRoot.length) {
      return aRecordRoot;
    }
  };

  const getAllValueWithCallbackRecursive = async (callback, name, type) => {

    // 1：先用 type 类型记录去查询
    const aValues = await callback(name, type);
    if (aValues && aValues.length) {
      return [aValues, name];
    }

    if (type === Packet.TYPE.NS) {
      return;
    }

    // 2：没记录的情况下，先查CNAME
    const aValues2 = await callback(name, Packet.TYPE.CNAME);
    if (!aValues2 || !aValues2.length) {
      return;
    }
    const sNewName = trimIllegalCharacters(randomArrayElement(aValues2));
    if (sNewName === name) {
      // 防止死循环
      return;
    }

    // 3：最后查询该新域名的 type 类型记录
    const aValues3 = await getAllValueWithCallbackRecursive(callback, sNewName, type);
    if (!aValues3) {
      return;
    }
    return aValues3;

  };

  const getAllValueByRecord = async (name, type, sDnsAddr) => {

    // 1：首先在本地数据库里递归查询域名解析记录
    const ret = await getAllValueWithCallbackRecursive(getAllRecordValueInDb, name, type);
    if (ret) {
      const [aValues, sNewName] = ret;
      return aValues;
    }

    // 2：都查不到才发送 DNS 请求查询
    const request = getPacket(name, type);
    if (!sDnsAddr) {
      if (type === Packet.TYPE.NS) {
        // 当 Db 无记录时，则查询 NS 记录必须提供 sDnsAddr
        return;
      }
      sDnsAddr = await getDnsAddr(name, iDnsAddrType);
    }
    if (!sDnsAddr) {
      return;
    }
    if (type !== Packet.TYPE.NS) {
      // 查询 NS 记录不支持 EDNS 选项
      pushEDNS(request);
    }
    const bufferData = await UDPClient(request.toBuffer(), sDnsAddr);
    const response = Packet.parse(bufferData);
    //console.log('=== UDPClient ===\r\n', response);
    const aNewZoneData = fromResponseToArrayZoneData(response);
    const mNewZoneData = fromZoneDataToMap(aNewZoneData);
    await db.fPutRecord('zone_cache', mNewZoneData);
    const aRecord3 = await db.fGetRecord('zone_cache', name, type);
    if (aRecord3) {
      return aRecord3;
    }

  };

  const getDnsAddrOne = async (name, type, sDnsAddr) => {
    // 1：先查NS记录
    const aNSvalue = await getAllValueByRecord(name, Packet.TYPE.NS, sDnsAddr);
    //console.log('getDnsAddrOne1 =>', name, aNSvalue);
    if (aNSvalue && aNSvalue.length) {
      // 2：查到NS记录后再去查服务器IP地址
      const sNSvalue = trimIllegalCharacters(randomArrayElement(aNSvalue));
      if (sNSvalue === name) {
        // 如果这个域名的NS记录就是自己则跳过，解决死循环
        return;
      }
      const aDnsAddr = await getAllValueByRecord(sNSvalue, type);
      if (aDnsAddr && aDnsAddr.length) {
        //console.log('getDnsAddrOne2 =>', name, sNSvalue, aDnsAddr);
        return randomArrayElement(aDnsAddr);
      }
    }
  };

  const getDnsAddr = async (name, type) => {
    // 根据域名查找DNS服务器的IP地址
    const names = name.split('.');
    const iArrayIndexMax = names.length - 1;
    var sDnsAddr;
    for (var i = 0; i <= iArrayIndexMax; i++) {
      const sName = names.slice(iArrayIndexMax - i, iArrayIndexMax + 1).join('.');
      //console.log('getDnsAddr', sName);
      const tmpAddr = await getDnsAddrOne(sName, type, sDnsAddr);
      //console.log('getDnsAddr', i, name, sName, tmpAddr, sDnsAddr);
      if (!tmpAddr) {
        return sDnsAddr;
      }
      sDnsAddr = tmpAddr
    }
    return sDnsAddr;
  };

  const fPutValueToAnswer = (answer, type, oValue) => {
    if (type === Packet.TYPE.CNAME) {
      answer.domain = oValue;
      return;
    }
    if (type === Packet.TYPE.MX) {
      answer.exchange = oValue[0];
      answer.priority = oValue[1];
      return;
    }
    if (type === Packet.TYPE.TXT) {
      answer.data = oValue;
      return;
    }
    if (type === Packet.TYPE.A || type === Packet.TYPE.AAAA) {
      answer.address = oValue;
      return;
    }
    answer.data = oValue;
    //answer.data = base64Decode(oValue).toString('utf8');

  };

  const fPutAnswer = (response, aValues, name, type) => {
    for (const oValue of aValues) {
      const answer = {
        name,
        type,
        class: Packet.CLASS.IN,
        ttl: 1,
      };
      fPutValueToAnswer(answer, type, oValue);
      response.answers.push(answer);
    }
    //console.log('cached', name, type, aValue);
  };

  const getResponseByRequest = async (request) => {
    const response = Packet.createResponseFromRequest(request);

    if (request.header.id === my_request_id) {
      // 如果是自身服务器发的查询则忽略，防止死循环。
      return response;
    }

    if (!request.questions) {
      return response;
    }

    const sQuery0 = request.questions[0];
    if (!sQuery0) {
      return response;
    }

    if (sQuery0.type === Packet.TYPE.PTR) {
      // 这种没必要解析的就直接返回
      return response;
    }

    // 开始对sQuery0进行解析
    const ret = await getAllValueWithCallbackRecursive(getAllValueByRecord, sQuery0.name, sQuery0.type);
    if (ret) {
      const [aValues, sNewName] = ret;
      fPutAnswer(response, aValues, sNewName, sQuery0.type);
      //console.log('=== new ===\r\n', response);
      return response;
    }

    return response;

  };

  const server = dns2.createServer({
    udp: true,
    handle: async (request, send, rinfo) => {
      send(await getResponseByRequest(request));
    }
  });

  server.on('request', (request, response, rinfo) => {
    console.log(request.header.id, request.questions);
  });

  server.on('requestError', (error) => {
    console.log('Client sent an invalid request', error);
  });

  server.on('listening', () => {
    console.log(server.addresses());
  });

  server.on('close', () => {
    console.log('server closed');
  });

  server.listen({
    // Optionally specify port, address and/or the family of socket() for udp server:
    udp: {
      port: 53,
      address: "",
      type: "udp4",  // IPv4 or IPv6 (Must be either "udp4" or "udp6")
    },
  });

};
main();


// eventually
//server.close();
